import pdfplumber
import csv
import re
# 搜索关键词：课时 1，课时 2，。。。
def generate_topic_list(num_topics):
    return ["课时{}".format(i + 1) for i in range(num_topics)]
    # return ["专题 {}".format(i + 1) for i in range(num_topics)]


def extract_text_to_table_and_add_bookmarks(pdf_path, num_topics):
    topic_list = generate_topic_list(num_topics)
    print(f"查找内容 '{topic_list}'")
    result_data = []
    with pdfplumber.open(pdf_path) as pdf:
        for topic in topic_list:
        # 遍历每个关键字，查找包含关键字的行内容及对应页码
            pattern = re.compile(r'\b{}\b'.format(topic))
            for page_num, page in enumerate(pdf.pages):
                text = page.extract_text()
                lines = text.splitlines()
                for line_index, line in enumerate(lines):
                    if re.search(pattern, line):
                        result_data.append((line, page_num + 1, topic))
            # 将包含关键字的行内容、对应页码、关键字组成元组添加到列表
                        print(f"找到内容 '{result_data}'")
    # 将结果保存到CSV表格中（方便查看提取的数据）
    with open('result12.csv', 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["内容", "页码", "关键字"])  # 写入表头
        for data in result_data:
            writer.writerow(data)



if __name__ == "__main__":
    
    pdf_path = r"E:\wow.pdf"
    
    num_topics = 3  # 这里可以根据实际需要的关键字数量进行修改
    extract_text_to_table_and_add_bookmarks(pdf_path, num_topics)